import os
import sys
import numpy as np
import math
from glob import glob
import argparse
from absl import app
from absl.flags import argparse_flags
from network import compress
from network import inference
from network import decompress

def encode(args):
    if not os.path.isdir(args.comp_dir):
        os.makedirs(args.comp_dir)
        
    test_file = args.test_file
    num = len(glob(test_file + '/*.png'))
    print("Number of pictures: {}".format(num))
    
    for image_file in glob(test_file + '/*.png'):
        print(image_file[(len(test_file)+1):])
        input = image_file
        output = args.comp_dir + "/" + image_file[(len(test_file)+1):-4] + '.npz'
        compress(input, output, args)  

def infer(args):
    if not os.path.isdir(args.infer_dir):
        os.makedirs(args.infer_dir)
        
    test_file = args.test_file
    num = len(glob(test_file + '/*.png'))
    print("Number of pictures: {}".format(num))
    
    for image_file in glob(test_file + '/*.png'):
        print(image_file[(len(test_file)+1):])
        input = image_file
        output = args.infer_dir + "/" + image_file[(len(test_file)+1):-4] + '.png'
        inference(input, output, args)  
        
def decode(args):
    if not os.path.isdir(args.rec_dir):
        os.makedirs(args.rec_dir)

    for image_file in glob(args.comp_dir + '/*.npz'):
        print(image_file[(len(args.comp_dir)+1):])
        input = image_file
        output = args.rec_dir + "/" + image_file[(len(args.comp_dir)+1):-4] + '.png'
        decompress(input, output, args)
  
def parse_args(argv):
    """Parses command line arguments."""
    parser = argparse_flags.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    
    subparsers = parser.add_subparsers(
        title="commands", dest="command",
        help="What to do: 'compress' reads an image file (lossless "
            "PNG format) and writes a compressed binary file. 'decompress' "
            "reads a binary file and reconstructs the image (in PNG format). ")
    
    # 'compress' subcommand.
    encode_cmd = subparsers.add_parser(
        "encode",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Reads a PNG file, compresses it, and writes a bin&npz file.")
    
    # 'inference' subcommand.
    infer_cmd = subparsers.add_parser(
        "infer",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Reads a png file, directly reconstructs the image, and do not "
                  "write a bin&npz file.")
    
    # 'decompress' subcommand.
    decode_cmd = subparsers.add_parser(
        "decode",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="Reads a bin&npz file, reconstructs the image, and writes back "
                  "a PNG file.")
    
    parser.add_argument(
        "--test_file", default="test_images",
        help="Glob pattern identifying teseting data. This pattern must expand "
            "to a list of RGB images in PNG format.")
    
    parser.add_argument(
        "--attention", "-a", action="store_true",
        help="Use attention mechanism.")
    
    parser.add_argument(
        "--num_filters", type=int, default=192,
        help="Number of filters per layer.")
    
    parser.add_argument(
        "--gpu","-g", type=str, default="0",
        help="The gpu index")
    
    parser.add_argument(
        "--checkpoint_dir", default="models/model.ckpt")
    
    parser.add_argument(
        "--comp_dir", default="comp_images",
        help="directory to save bin&npz file using encode_cmd")
    
    parser.add_argument(
        "--infer_dir", default="infer_images",
        help="directory to save png file using infer_cmd")
    
    parser.add_argument(
        "--rec_dir", default="rec_images",
        help="directory to save png file using decode_cmd")

    args = parser.parse_args(argv[1:])

    if args.command is None:
        parser.print_usage()
        sys.exit(2)
    return args

def main(args):
    # Invoke subcommand.
    os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
    if args.command == "encode":
        encode(args)
    elif args.command == "infer":
        infer(args)
    elif args.command == "decode":
        decode(args)

if __name__ == "__main__":  
    app.run(main, flags_parser=parse_args)
